# -*- coding: utf-8 -*-
"""
# --------------------------------------------------------
# @Project: python-learning-notes
# @Author : panjq
# @E-mail : pan_jinquan@163.com
# @Date   : 2020-07-22 19:40:49
# --------------------------------------------------------
"""
from typing import List

import cv2
import copy
import numpy as np
import numbers
import torch
import PIL.Image as Image
import random
import math
from detection.transforms import affine_transform

"""
pct_thresh: 0.01表示如果长宽有一个不足图像的1%，则去除
iou_thresh: clip会缩小boxes，因此如果clip前后的IOU小于该阈值，box会被丢弃
"""
pct_thresh = 0.005
iou_thresh = 0.1
color = (127, 127, 127)


class Compose(object):
    """Composes several augmentations together.
    Args:
        transforms (List[Transform]): list of transforms to compose.
    Example:
    augmentations.Compose([
                    transforms.CenterCrop(10),
                    transforms.ToTensor(),
    ])
    """

    def __init__(self, transforms):
        self.transforms = [t for t in transforms if t]

    def __call__(self, image, boxes=None, labels=None):
        for t in self.transforms:
            image, boxes, labels = t(image, boxes, labels)
        return image, boxes, labels


class RandomChoice(object):
    """Apply single transformation randomly picked from a list. This transform does not support torchscript.
    """

    def __init__(self, transforms):
        if not isinstance(transforms, list):
            raise TypeError("Argument transforms should be a sequence")
        self.transforms = transforms

    def __call__(self, image, boxes=None, labels=None):
        t = random.choice(self.transforms)
        return t(image, boxes, labels)


def scale_size(size, scale=1.0):
    """保证size的结果是偶数"""
    if isinstance(scale, numbers.Number):
        scale = [scale, scale]
    s0 = int(size[0] * scale[0])
    s1 = int(size[1] * scale[1])
    s0 = s0 if s0 % 2 == 0 else s0 - 1
    s1 = s1 if s1 % 2 == 0 else s1 - 1
    dsize = [s0, s1]
    return dsize


class ToTensor(object):
    def __call__(self, img, boxes=None, labels=None, **kwargs):
        """
        :param img:
        :param boxes:
        :param labels:
        :return:
        """
        # [H,W,C]->[C,H,W]: (360, 480, 3)->(3, 360, 480)
        # img = img.astype(np.float32)
        # boxes = boxes.astype(np.float32)
        img = torch.from_numpy(img).permute(2, 0, 1)
        if kwargs:
            return img, boxes, labels, kwargs
        return img, boxes, labels


class TargetTransforms():
    def __init__(self, target="ssd"):
        """
        :param target:
        """
        self.target = target

    def yolo_target(self, image, boxes=None, labels=None):
        # （label,xmin,ymin,xmax,ymax）
        if len(boxes) == 0:
            target = np.empty(shape=(0, 5), dtype=np.float32)
        else:
            target = np.concatenate([labels.reshape(-1, 1), boxes], axis=1)
            target = np.asarray(target, dtype=np.float32)
        return image, target

    def ssd_target(self, image, boxes=None, labels=None):
        # （xmin,ymin,xmax,ymax,label）
        if len(boxes) == 0:
            target = np.empty(shape=(0, 5), dtype=np.float32)
        else:
            target = np.concatenate([boxes, labels.reshape(-1, 1)], axis=1)
            target = np.asarray(target, dtype=np.float32)
        return image, target

    def __call__(self, image, boxes=None, labels=None):
        if self.target == "ssd":
            image, target = self.ssd_target(image, boxes, labels)
        else:
            image, target = self.yolo_target(image, boxes, labels)
        return target


def check_bboxes(boxes):
    """
    :param boxes:
    :return:
    """
    for b in boxes:
        xmin, ymin, xmax, ymax = b
        assert xmax > xmin
        assert ymax > ymin


class ProtectBoxes(object):
    def __init__(self, norm=False, pct_th=pct_thresh, iou_th=iou_thresh):
        """
        限制boxes边界范围，防止越界的情况，去除小框，避免出现Nan值
        :param norm: True：输入的boxes是归一化坐标，即boxes/[w,h,w,h]，范围(0,1)
                    False：输入的boxes是像素坐标，即(xmin,ymin,xmax,ymax)
        :param pct_th: 0.01表示如果长宽有一个不足图像的1%，则去除
        :param iou_th: clip会缩小boxes，因此如果clip前后的IOU小于该阈值，box会被丢弃
        """
        # 用于限制boxes边界范围，防止越界的情况
        self.clip = ClipBoxes(norm=norm, iou_th=iou_th)
        # 去除小框，避免出现Nan值
        self.min_boxes = IgnoreMinBoxes(norm=norm, pct_th=pct_th)

    def __call__(self, image, boxes=None, labels=None):
        image, boxes, labels = self.clip(image, boxes, labels)
        image, boxes, labels = self.min_boxes(image, boxes, labels)
        return image, boxes, labels


class ClipBoxes(object):
    def __init__(self, norm=False, iou_th=iou_thresh):
        """
        用于限制boxes边界范围，防止越界的情况
        :param norm: True：输入的boxes是归一化坐标，即boxes/[w,h,w,h]，范围(0,1)
                    False：输入的boxes是像素坐标，即(xmin,ymin,xmax,ymax)
        :param iou_th: clip会缩小boxes，因此如果clip前后的IOU小于该阈值，box会被丢弃
        """
        self.norm = norm
        self.iou_th = iou_th
        self.index = []

    @staticmethod
    def cal_iou(box1, box2):
        """
        :param box1: = [xmin1, ymin1, xmax1, ymax1]
        :param box2: = [xmin2, ymin2, xmax2, ymax2]
        :return:
        """
        xmin1, ymin1, xmax1, ymax1 = box1[:4]
        xmin2, ymin2, xmax2, ymax2 = box2[:4]
        # 计算每个矩形的面积
        s1 = (xmax1 - xmin1) * (ymax1 - ymin1)  # C的面积
        s2 = (xmax2 - xmin2) * (ymax2 - ymin2)  # G的面积

        # 计算相交矩形
        xmin = max(xmin1, xmin2)
        ymin = max(ymin1, ymin2)
        xmax = min(xmax1, xmax2)
        ymax = min(ymax1, ymax2)

        w = max(0, xmax - xmin)
        h = max(0, ymax - ymin)
        area = w * h  # C∩G的面积
        iou = area / (s1 + s2 - area)
        return iou

    def check_overlap(self, image, boxes, labels, cboxes, iou_th):
        """
        检测clip前后的boxes的IOU变化，因此如果clip前后的IOU小于该阈值iou_th，box会被丢弃
        :param image:
        :param boxes:clip前的boxes
        :param labels:
        :param cboxes:clip后的boxes
        :param iou_th:
        :return:
        """
        ious = np.zeros(shape=(len(boxes),))
        for i in range(len(boxes)):
            ious[i] = self.cal_iou(boxes[i], cboxes[i])
        # print("ious:{}".format(ious))
        # print("labels:{}".format(labels))
        index = ious > iou_th
        boxes = cboxes[index]
        labels = labels[index]
        self.index = index
        return image, boxes, labels

    def get_index(self):
        return self.index

    def clip(self, image, boxes, labels):
        """限制boxes边界范围"""
        cboxes = boxes.copy()
        if self.norm:
            cboxes = np.clip(boxes, 0, 1)
        else:
            height, width, _ = image.shape
            cboxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, width - 1)
            cboxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, height - 1)
        return image, cboxes, labels

    def __call__(self, image, boxes=None, labels=None):
        if len(boxes) > 0:
            image, cboxes, labels = self.clip(image, boxes, labels)
            image, boxes, labels = self.check_overlap(image, boxes, labels, cboxes, iou_th=self.iou_th)
        check_bboxes(boxes)
        return image, boxes, labels


class IgnoreMinBoxes(object):
    def __init__(self, norm=True, pct_th=pct_thresh):
        """
        去除小框，避免出现Nan值
        :param norm: True：输入的boxes是归一化坐标，即boxes/[w,h,w,h]，范围(0,1)
                    False：输入的boxes是像素坐标，即(xmin,ymin,xmax,ymax)
        :param pct_th: 0.01表示如果长宽有一个不足图像的1%，则去除
        """
        self.norm = norm
        self.pct_th = pct_th

    def __call__(self, image, boxes=None, labels=None):
        if len(boxes) > 0:
            height, width, _ = image.shape
            if not self.norm:
                boxes = boxes / [width, height, width, height]
            w = boxes[:, 2] - boxes[:, 0]
            h = boxes[:, 3] - boxes[:, 1]
            index0 = w > self.pct_th
            index1 = h > self.pct_th
            index = index0 * index1
            boxes = boxes[index, :]
            labels = labels[index]
            if not self.norm:
                boxes = boxes * [width, height, width, height]
        return image, boxes, labels


class IgnoreBadBoxes(object):
    """去除不合法的bbox,如小于0，越界的bbox"""

    def __init__(self):
        pass

    def get_mask_index(self, mask):
        """
        :param mask:
        :return:
        """
        index = [sum(mask[i, :]) == len(mask[i, :]) for i in range(len(mask))]
        index = np.array(index)
        return index

    def __call__(self, image, boxes=None, labels=None):
        if len(boxes) > 0:
            height, width, _ = image.shape
            index0 = self.get_mask_index(boxes >= 0)
            index1 = self.get_mask_index(boxes < [width, height, width, height])
            index = index0 * index1
            boxes = boxes[index]
            labels = labels[index]
        check_bboxes(boxes)
        return image, boxes, labels


class Normalize(object):
    """归一化数据"""

    def __init__(self, mean=0.0, std=1.0, norm=True):
        """
        :param mean:
        :param std:
        :param norm: 0-1 Normalize(True)
        """
        self.mean = mean
        self.std = std
        self.norm = norm

    def __call__(self, image, boxes=None, labels=None):
        image = image.astype(np.float32)
        if self.norm:
            image /= 255.0
        image -= self.mean
        image /= self.std
        # check_bboxes(boxes)
        return image, boxes, labels


class UnNormalizeBoxesCoords(object):
    """将归一化boxes坐标转换为图像坐标"""

    def __call__(self, image, boxes=None, labels=None):
        if len(boxes) > 0:
            height, width, channels = image.shape
            boxes = boxes * [width, height, width, height]
        boxes = np.asarray(boxes, np.float32)
        return image, boxes, labels


class NormalizeBoxesCoords(object):
    """将图像坐标转为归一化boxes坐标"""

    def __call__(self, image, boxes=None, labels=None):
        if len(boxes) > 0:
            height, width, channels = image.shape
            boxes = boxes / [width, height, width, height]
        boxes = np.asarray(boxes, np.float32)
        return image, boxes, labels


class YOLOCoords(object):
    def __init__(self, max_boxes=120, norm=False):
        """
        转换为YOLO格式
        :param max_boxes 最大boxes的个数，不足用全0填充
        :param norm 坐标是否进行坐标归一化： x/image_width,y/image_height
        """
        self.max_boxes = max_boxes
        self.norm = norm

    @staticmethod
    def xyxy2cxcywh(xyxy: np.ndarray, width, height, normalized=False):
        """(xmin, ymin, xmax, ymax)==>(cx,cy,w,h)"""
        cxcywh = np.zeros_like(xyxy)
        cxcywh[:, 0] = (xyxy[:, 2] + xyxy[:, 0]) / 2  # cx
        cxcywh[:, 1] = (xyxy[:, 3] + xyxy[:, 1]) / 2  # cy
        cxcywh[:, 2] = (xyxy[:, 2] - xyxy[:, 0])  # w
        cxcywh[:, 3] = (xyxy[:, 3] - xyxy[:, 1])  # h
        if normalized:
            cxcywh = cxcywh / (width, height, width, height)
        return cxcywh

    def __call__(self, image, boxes, labels):
        dboxes = np.zeros((self.max_boxes, 4))
        dlabels = np.zeros((self.max_boxes))
        if len(boxes) > 0:
            width, height, _ = image.shape
            boxes = self.xyxy2cxcywh(boxes, width, height, normalized=self.norm)
            # Fix a Bug: len(cboxes) > self.max_boxes
            if len(boxes) > self.max_boxes:
                boxes = boxes[0:self.max_boxes]
                labels = labels[0:self.max_boxes]
            dboxes[0:len(boxes)] = boxes
            dlabels[0:len(boxes)] = labels
        return image, dboxes, dlabels


class Resize(object):
    """resize"""

    def __init__(self, size=[300, 300]):
        """
        图像resize,可能会导致图像内容失真
        :param size:
        """
        self.size = tuple(size)

    def __call__(self, image, boxes=None, labels=None):
        height, width, _ = image.shape
        if not boxes is None and len(boxes) > 0:
            scale = [self.size[0] / width, self.size[1] / height] * 2
            boxes = boxes * scale
        image = cv2.resize(image, (self.size[0], self.size[1]))
        return image, boxes, labels


class ResizePadding(object):
    def __init__(self, size=[300, 300], value=color):
        """
        等比例图像resize,保持原始图像内容比，避免失真,短边会0填充
        :param size:
        """
        self.size = tuple(size)
        self.value = tuple(value)

    def get_pad_size(self, image_shape, size):
        height, width, _ = image_shape
        scale = min([size[0] / width, size[1] / height])
        new_size = [int(width * scale), int(height * scale)]
        pad_w = size[0] - new_size[0]
        pad_h = size[1] - new_size[1]
        top, bottom = pad_h // 2, pad_h - (pad_h // 2)
        left, right = pad_w // 2, pad_w - (pad_w // 2)
        pad_size = [top, bottom, left, right]
        return new_size, pad_size, scale

    def __call__(self, image, boxes=None, labels=None):
        new_size, pad_size, scale = self.get_pad_size(image.shape, size=self.size)
        top, bottom, left, right = pad_size
        if not boxes is None and len(boxes) > 0:
            boxes = boxes * scale
            boxes = boxes + [left, top, left, top]
        image = cv2.resize(image, (new_size[0], new_size[1]))
        image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, None, value=self.value)
        return image, boxes, labels

    def inverse(self, image, boxes=None, labels=None):
        new_size, pad_size, scale = self.get_pad_size(image.shape, size=self.size)
        top, bottom, left, right = pad_size
        if not boxes is None and len(boxes) > 0:
            boxes = boxes - [left, top, left, top]
            boxes = boxes / scale
        return image, boxes, labels


class ScaleResize(object):
    """缩放"""

    def __init__(self, size, scale: float = 1.0, padding: bool = False):
        """
        :param size:
        :param scale:
        :param padding:
        :param interpolation:
        """
        dsize = scale_size(size, scale)
        if padding:
            self.resizer = ResizePadding(size=dsize)
        else:
            self.resizer = Resize(size=dsize)

    def __call__(self, image, boxes, labels):
        image, boxes, labels = self.resizer(image, boxes, labels)
        return image, boxes, labels


class MultiScaleResize(object):
    """多尺度缩放"""

    def __init__(self, size, scale: List[float], padding: bool = False):
        """
        :param size:
        :param scale:
        :param padding:
        """
        resizer = []
        for s in scale:
            resize = ScaleResize(size, s, padding=padding)
            resizer.append(resize)
        self.resizer = RandomChoice(resizer)

    def __call__(self, image, boxes, labels):
        image, boxes, labels = self.resizer(image, boxes, labels)
        return image, boxes, labels


class MultiScaleRandomCrop(object):
    """多尺度缩放-->transform-->随机裁剪"""

    def __init__(self, size, scale: List[float], padding: bool = False, transform=None):
        """
        :param size: 随机裁剪输出的结果
        :param scale: 在size的基础身上进行多尺度缩放
        :param padding:
        :param transform:
        """
        resizer = []
        for s in scale:
            resize = ScaleResize(size, s, padding=padding)
            resizer.append(resize)
        self.resizer = RandomChoice(resizer)
        self.crop = RandomCropLarge(size=size, p=1.0)
        self.transform = transform
        if isinstance(self.transform, list):
            self.transform = Compose(self.transform)

    def __call__(self, image, boxes, labels):
        image, boxes, labels = self.resizer(image, boxes, labels)
        if self.transform:
            image, boxes, labels = self.transform(image, boxes, labels)
        image, boxes, labels = self.crop(image, boxes, labels)
        return image, boxes, labels


class RandomResizePadding():
    def __init__(self, size, p=0.5):
        """
        随机使用ResizePadding和Resize，以提高泛化性
        :param p: 随机概率
        """
        self.p = p
        self.resize_padding = ResizePadding(size)
        self.resize = Resize(size)

    def __call__(self, img, boxes, labels):
        if random.random() < self.p:
            img, boxes, labels = self.resize_padding(img, boxes, labels)
        else:
            img, boxes, labels = self.resize(img, boxes, labels)
        return img, boxes, labels


class RandomAffineResizePadding(object):
    def __init__(self, degrees=15, output_size=None, p=0.5, check=True):
        """
        保持原始图像内容比，避免失真,短边会0填充，随机旋转进行仿生变换
        PS：如果训练时加入“保持原始图像内容比”的数据增强，那测试也建议加上“保持原始图像内容比”
        不然测试效果会差一点
        :param degrees:随机旋转的角度
        :param output_size: 仿生变换输出的大小
        :param p: 仿生变换的概率P:
        :param check: True False，检测使用ToPercentCoords()或ProtectBBoxes限制box越界情况
        """
        if isinstance(degrees, numbers.Number):
            if degrees < 0:
                raise ValueError("If degrees is a single number, it must be positive.")
            self.degrees = (-degrees, degrees)
        else:
            if len(degrees) != 2:
                raise ValueError("If degrees is a sequence, it must be of len 2.")
            self.degrees = degrees
        self.output_size = output_size
        # 去除因为旋转导致越界的boxes
        self.check = check
        self.p = p
        self.protect_bboxes = ProtectBoxes(norm=False)
        if self.output_size:
            self.size = Resize(self.output_size)

    def affine(self, image, boxes, labels, angle, output_size):
        if not output_size:
            h, w, _ = image.shape
            output_size = [w, h]
        image, boxes, center, scale, kwargs = affine_transform.AffineTransform. \
            affine_transform(image, boxes, output_size, rot=angle)
        return image, boxes, labels

    def __call__(self, image, boxes, labels):
        '''
        :param image: nparray img
        :param boxes: np.array([[88, 176, 250, 312, 1222], [454, 115, 500, 291, 1222]]), 里面为x1, y1, x2, y2, 标签
        :return:
        '''
        if random.random() < self.p and len(boxes) > 0:
            angle = random.uniform(self.degrees[0], self.degrees[1])
            image, boxes, labels = self.affine(image, boxes, labels, angle, self.output_size)
        elif self.output_size:
            image, boxes, labels = self.size(image, boxes, labels)
        # 去除因为旋转导致越界的boxes
        if self.check:
            image, boxes, labels = self.protect_bboxes(image, boxes, labels)
        return image, boxes, labels

    def set_degrees(self, degrees):
        self.degrees = [degrees, degrees]


class RandomBoxNoise(object):
    def __init__(self, p=0.5, noise=[3, 3]):
        """
        boxes添加随机扰动
        :param p:
        :param noise: [noise_x,noise_y]添加(x,y)方向的扰动，noise值越大，Box扰动越大
        """
        self.noise = noise
        self.p = p

    def __call__(self, image, boxes=None, labels=None):
        # boxes = np.asarray(boxes)
        if random.random() < self.p:
            rx = int(random.uniform(-self.noise[0], self.noise[0]))
            ry = int(random.uniform(-self.noise[1], self.noise[1]))
            boxes = boxes + [rx, ry] * 2
            height, width, _ = image.shape
            boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, width - 1)
            boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, height - 1)
        return image, boxes, labels


class RandomExpand(object):
    def __init__(self, ratio=1.2, padding=0, p=0.5):
        self.ratio = ratio
        self.p = p
        self.padding = padding

    def __call__(self, image, boxes, labels):
        if random.random() < self.p:
            height, width, depth = image.shape
            ratio = random.uniform(1, self.ratio)
            left = random.uniform(0, width * ratio - width)
            top = random.uniform(0, height * ratio - height)
            expand_image = np.zeros((int(height * ratio), int(width * ratio), depth), dtype=image.dtype)
            expand_image[:, :, :] = self.padding
            expand_image[int(top):int(top + height), int(left):int(left + width)] = image
            image = expand_image
            if len(boxes) > 0:
                boxes[:, :2] += (int(left), int(top))
                boxes[:, 2:] += (int(left), int(top))
        return image, boxes, labels


class RandomCropLarge(object):
    def __init__(self, size=[320, 320], p=0.8):
        """
        RandomCropLarge与RandomCrop类似都是实现随机裁剪，
        RandomCrop会自定计算裁剪区域，为保证不会裁剪掉boxes的区域，其裁剪幅度较小
        RandomCropLarge可设定最小裁剪区域size，裁剪幅度较大，可能会裁剪部分box区域
        PS： RandomCropLarge建议放在ResizePadding和Resize之前
        :param size: 最小crop的大小[W,H],注意，如果输入的图像与size一样大小，相当于没有裁剪
        :param p:概率P
        :param iou_th: clip会缩小boxes，因此如果clip前后的IOU小于该阈值，box会被丢弃
        """
        self.p = p
        self.size = size
        self.clip = ClipBoxes(norm=False)

    def __call__(self, image, boxes, labels):
        if random.random() < self.p:
            return self.random_crop(image, boxes, labels)
        return image, boxes, labels

    def random_box(self, image, size):
        """
        Get parameters for ``crop`` for a random crop.
        :param image:
        :param size: 最小crop的大小
        :return:
        """
        h, w, d = image.shape
        # tw, th = int(random.uniform(size[0], w)), int(random.uniform(size[1], h))
        tw, th = min(size[0], w), min(size[1], h)
        y = int(random.uniform(0, h - th))
        x = int(random.uniform(0, w - tw))
        box = [x, y, x + tw, y + th]
        return box

    def random_crop(self, image, boxes, labels):
        for i in range(20):
            box = self.random_box(image, self.size)
            cimage = image[box[1]: box[3], box[0]: box[2]]
            cboxes = boxes.copy()
            clabels = labels.copy()
            cboxes[:, 0::2] = cboxes[:, 0::2] - box[0]
            cboxes[:, 1::2] = cboxes[:, 1::2] - box[1]
            cimage, cboxes, clabels = self.clip(cimage, cboxes, clabels)
            # 为了避免因裁剪导致boxes为空，当len(boxes)=0时继续循环，最多循环100次
            if len(cboxes) > 0:
                return cimage, cboxes, clabels
        return image, boxes, labels


class RandomCrop(object):
    """ 随机裁剪"""

    def __init__(self, p=0.5, margin_rate=0.5):
        """
        RandomCropLarge与RandomCrop类似都是实现随机裁剪，
        RandomCrop会自定计算裁剪区域，为保证不会裁剪掉boxes的区域，其裁剪幅度较小
        RandomCropLarge可设定最小裁剪区域size，裁剪幅度较大，可能会裁剪部分box区域
        :param p: 实现随机裁剪的概率
        :param margin_rate: 随机裁剪的幅度
        """
        self.p = p
        self.margin_rate = margin_rate

    def __call__(self, img, boxes, labels):
        """
        Args:
            img (PIL Image): Image to be cropped.

        Returns:
            PIL Image: Cropped image.
        """

        if random.random() < self.p:
            h_img, w_img, _ = img.shape
            # 得到可以包含所有bbox的最大bbox
            if len(boxes) > 0:
                max_bbox = np.concatenate([np.min(boxes[:, 0:2], axis=0), np.max(boxes[:, 2:4], axis=0)], axis=-1)
            else:
                max_bbox = [int(w_img * 0.2), int(h_img * 0.2), int(w_img * 0.8), int(h_img * 0.8)]
            max_l_trans = max_bbox[0]
            max_u_trans = max_bbox[1]
            max_r_trans = w_img - max_bbox[2]
            max_d_trans = h_img - max_bbox[3]
            # crop_xmin = max(0, int(max_bbox[0] - random.uniform(0, max_l_trans)))
            # crop_ymin = max(0, int(max_bbox[1] - random.uniform(0, max_u_trans)))
            # crop_xmax = max(w_img, int(max_bbox[2] + random.uniform(0, max_r_trans)))
            # crop_ymax = max(h_img, int(max_bbox[3] + random.uniform(0, max_d_trans)))
            crop_xmin = int(max(0, random.uniform(0, max_l_trans * self.margin_rate)))
            crop_ymin = int(max(0, random.uniform(0, max_u_trans * self.margin_rate)))
            crop_xmax = int(min(w_img, random.uniform(w_img - max_r_trans * self.margin_rate, w_img)))
            crop_ymax = int(min(h_img, random.uniform(h_img - max_d_trans * self.margin_rate, h_img)))

            img = img[crop_ymin: crop_ymax, crop_xmin: crop_xmax]
            h_img, w_img, _ = img.shape
            if len(boxes) > 0:
                # boxes[:, [0, 2]] = boxes[:, [0, 2]] - crop_xmin
                # boxes[:, [1, 3]] = boxes[:, [1, 3]] - crop_ymin
                boxes = boxes - [crop_xmin, crop_ymin, crop_xmin, crop_ymin]
                boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, w_img - 1)
                boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, h_img - 1)
        return img, boxes, labels


class RandomAffine(object):
    """ 随机Padding"""

    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, img, boxes, labels):
        if random.random() < self.p:
            h_img, w_img, _ = img.shape
            # 得到可以包含所有bbox的最大bbox
            max_bbox = np.concatenate([np.min(boxes[:, 0:2], axis=0), np.max(boxes[:, 2:4], axis=0), ], axis=-1)
            max_l_trans = max_bbox[0]
            max_u_trans = max_bbox[1]
            max_r_trans = w_img - max_bbox[2]
            max_d_trans = h_img - max_bbox[3]
            tx = random.uniform(-(max_l_trans - 1), (max_r_trans - 1))
            ty = random.uniform(-(max_u_trans - 1), (max_d_trans - 1))

            M = np.array([[1, 0, tx], [0, 1, ty]])
            img = cv2.warpAffine(img, M, (w_img, h_img))
            boxes[:, [0, 2]] = boxes[:, [0, 2]] + tx
            boxes[:, [1, 3]] = boxes[:, [1, 3]] + ty
        return img, boxes, labels


class RandomRotation(object):
    def __init__(self, degrees=15, p=0.5, check=True, value=color):
        """
        随机旋转，BUG,loss出现inf,可能因为旋转导致部分box越界或者丢失
        :param degrees:
        :param p:随机旋转的概率
        :param check: True False:
        """
        if isinstance(degrees, numbers.Number):
            if degrees < 0:
                raise ValueError("If degrees is a single number, it must be positive.")
            self.degrees = (-degrees, degrees)
        else:
            if len(degrees) != 2:
                raise ValueError("If degrees is a sequence, it must be of len 2.")
            self.degrees = degrees
        self.p = p
        # 去除因为旋转导致越界的boxes
        self.check = check
        self.value = value
        self.protect_bboxes = ProtectBoxes(norm=False)

    def rotation_v1(self, image, boxes, labels, angle):
        """
        :param image:
        :param boxes:
        :param labels:
        :param angle: 旋转角度
        :return:
        """
        h, w, _ = image.shape
        center = (w / 2., h / 2.)
        rot_mat = cv2.getRotationMatrix2D(center, angle, 1.0)
        image = cv2.warpAffine(image, rot_mat, dsize=(w, h), borderValue=self.value)
        num_boxes = len(boxes)
        if num_boxes > 0:
            points, num_boxes = affine_transform.get_boxes2points(boxes)
            for i in range(num_boxes):
                points[i, :] = affine_transform.rotate_points(points[i, :], centers=[center], angle=angle, height=h)
            boxes = affine_transform.get_points2bboxes(points)
        return image, boxes, labels

    def rotation_v2(self, image, boxes, labels, angle, output_size=None):
        """
        :param image:
        :param boxes:
        :param labels:
        :param angle: 旋转角度
        :param output_size: 输出大小
        :return:
        """
        h, w, _ = image.shape
        if not output_size:
            output_size = [w, h]
        image, boxes = affine_transform.affine_transform_for_boxes(image,
                                                                   boxes,
                                                                   output_size=output_size,
                                                                   rot=angle,
                                                                   value=self.value)
        return image, boxes, labels

    def __call__(self, image, boxes, labels):
        '''
        :param image: nparray img
        :param boxes: np.array([[88, 176, 250, 312, 1222], [454, 115, 500, 291, 1222]]), 里面为x1, y1, x2, y2, 标签
        :param p: 随机比例
        :return:
        '''
        if random.random() < self.p and len(boxes) > 0:
            angle = random.uniform(self.degrees[0], self.degrees[1])
            image, boxes, labels = self.rotation_v1(image, boxes, labels, angle)
            # 去除因为旋转导致越界的boxes
            if self.check:
                image, boxes, labels = self.protect_bboxes(image, boxes, labels)
        return image, boxes, labels

    def set_degrees(self, degrees):
        self.degrees = [degrees, degrees]


class RandomRot90(object):
    """ 随机旋转90度"""

    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, img, boxes, labels):
        '''
        :param img: nparray img
        :param boxes: np.array([[88, 176, 250, 312, 1222], [454, 115, 500, 291, 1222]]), 里面为x1, y1, x2, y2, 标签
        :param p: 随机比例
        :return:
        '''
        # 顺时针旋转90度
        if random.random() < self.p:
            h, w, _ = img.shape
            img = cv2.transpose(img)
            img = cv2.flip(img, 1)
            if len(boxes) > 0:
                # bounding crop_type 的变换: 一个图像的宽高是W,H, 如果顺时90度转换，那么原来的原点(0, 0)到了 (H, 0) 这个最右边的顶点了，
                # 设图像中任何一个转换前的点(x1, y1), 转换后，x1, y1是表示到 (H, 0)这个点的距离，所以我们只要转换回到(0, 0) 这个点的距离即可！
                # 所以+90度转换后的点为 (H-y1, x1), -90度转换后的点为(y1, W-x1)
                boxes[:, [0, 1, 2, 3]] = boxes[:, [1, 0, 3, 2]]
                boxes[:, [0, 2]] = h - boxes[:, [2, 0]]
        return img, boxes, labels


class RandomHorizontalFlip(object):
    """Horizontally flip the given PIL Image randomly with a given probability.
    随机水平翻转
    Args:
        p (float): probability of the image being flipped. Default value is 0.5
    """

    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, image, boxes, labels):
        """
        :param image:
        :param boxes: xmin,ymin,xmax,ymax
        :param labels:
        :return:
        """
        if random.random() < self.p:
            height, width, _ = image.shape
            image = image[:, ::-1, :]
            if len(boxes) > 0:
                boxes[:, [0, 2]] = width - boxes[:, [2, 0]]
        return image, boxes, labels

    def __repr__(self):
        return self.__class__.__name__ + '(p={})'.format(self.p)


class RandomVerticalFlip(object):
    """Vertically flip the given PIL Image randomly with a given probability.
    随机上下翻转
    Args:
        p (float): probability of the image being flipped. Default value is 0.5
    """

    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, image, boxes, classes):
        """
        :param image:
        :param boxes: xmin,ymin,xmax,ymax
        :param classes:
        :return:
        """
        if random.random() < self.p:
            height, width, _ = image.shape
            image = image[::-1, :, :]
            if len(boxes) > 0:
                boxes[:, [1, 3]] = height - boxes[:, [3, 1]]
        return image, boxes, classes

    def __repr__(self):
        return self.__class__.__name__ + '(p={})'.format(self.p)


class SwapChannels(object):
    """交换图像颜色通道的顺序"""

    def __init__(self, swaps=[], p=1.0):
        """
        由于输入可能是RGB或者BGR格式的图像，随机交换通道顺序可以避免图像通道顺序的影响
        :param swaps:指定交换的颜色通道顺序，如[2,1,0]
                     如果swaps=[]或None，表示随机交换顺序
        :param p:概率
        """
        self.p = p
        self.swap_list = []
        if not swaps:
            self.swap_list = [[0, 1, 2], [2, 1, 0]]
        else:
            self.swap_list = [swaps]
        self.swap_index = np.arange(len(self.swap_list))

    def __call__(self, image, boxes=None, labels=None):
        if random.random() < self.p:
            index = np.random.choice(self.swap_index)
            swap = self.swap_list[index]
            image = image[:, :, swap]
        return image, boxes, labels


class RandomSaturation(object):
    def __init__(self, p=0.5, lower=0.5, upper=1.5):
        self.lower = lower
        self.upper = upper
        self.p = p
        assert self.upper >= self.lower, "contrast upper must be >= lower."
        assert self.lower >= 0, "contrast lower must be non-negative."

    def __call__(self, image, boxes=None, labels=None):
        if random.random() < self.p:
            image[:, :, 1] *= random.uniform(self.lower, self.upper)

        return image, boxes, labels


class RandomHue(object):
    def __init__(self, delta=18.0):
        assert delta >= 0.0 and delta <= 360.0
        self.delta = delta

    def __call__(self, image, boxes=None, labels=None):
        if random.randint(2):
            image[:, :, 0] += random.uniform(-self.delta, self.delta)
            image[:, :, 0][image[:, :, 0] > 360.0] -= 360.0
            image[:, :, 0][image[:, :, 0] < 0.0] += 360.0
        return image, boxes, labels


class RandomLightingNoise(object):
    def __init__(self):
        self.perms = ((0, 1, 2), (0, 2, 1),
                      (1, 0, 2), (1, 2, 0),
                      (2, 0, 1), (2, 1, 0))

    def __call__(self, image, boxes=None, labels=None):
        if random.randint(2):
            swap = self.perms[random.randint(len(self.perms))]
            shuffle = SwapChannels(swap)  # shuffle channels
            image = shuffle(image)
        return image, boxes, labels


class ConvertColor(object):
    def __init__(self, current, transform):
        self.transform = transform
        self.current = current

    def __call__(self, image, boxes=None, labels=None):
        if self.current == 'BGR' and self.transform == 'HSV':
            image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        elif self.current == 'RGB' and self.transform == 'HSV':
            image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
        elif self.current == 'BGR' and self.transform == 'RGB':
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        elif self.current == 'HSV' and self.transform == 'BGR':
            image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
        elif self.current == 'HSV' and self.transform == "RGB":
            image = cv2.cvtColor(image, cv2.COLOR_HSV2RGB)
        else:
            raise NotImplementedError
        return image, boxes, labels


class RandomColorJitter(object):
    def __init__(self, p=0.5, brightness=0.5, contrast=0.5, saturation=0.1, hue=0.1):
        """
        :param p:
        :param brightness (float or tuple of float (min, max)): How much to jitter brightness.
            brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness]
            or the given [min, max]. Should be non negative numbers.
        :param contrast (float or tuple of float (min, max)): How much to jitter contrast.
            contrast_factor is chosen uniformly from [max(0, 1 - contrast), 1 + contrast]
            or the given [min, max]. Should be non negative numbers.
        :param saturation (float or tuple of float (min, max)): How much to jitter saturation.
            saturation_factor is chosen uniformly from [max(0, 1 - saturation), 1 + saturation]
            or the given [min, max]. Should be non negative numbers.
        :param hue (float or tuple of float (min, max)): How much to jitter hue.
            hue_factor is chosen uniformly from [-hue, hue] or the given [min, max].
            Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5.(色调建议设置0.1，避免颜色变化过大)
        """
        from torchvision import transforms
        self.p = p
        self.random_choice = transforms.RandomChoice([
            transforms.ColorJitter(brightness=brightness),
            transforms.ColorJitter(contrast=contrast),
            transforms.ColorJitter(saturation=saturation),
            transforms.ColorJitter(hue=hue),
        ])
        self.color_transforms = transforms.ColorJitter(brightness=brightness,
                                                       contrast=contrast,
                                                       saturation=saturation,
                                                       hue=hue)

    def __call__(self, image, boxes=None, labels=None):
        if random.random() < self.p:
            is_pil = isinstance(image, np.ndarray)
            if is_pil:
                image = Image.fromarray(image)
            # image = self.random_choice(image)
            image = self.color_transforms(image)
            if is_pil:
                image = np.asarray(image)
        return image, boxes, labels


class RandomContrastBrightness(object):
    def __init__(self, p=0.5, lower=0.5, upper=1.5, delta=32):
        """
        随机调整对比度RandomContrast和亮度RandomBrightness (概率：0.5)
        :param p:
        :param lower:
        :param upper:
        """
        self.p = p
        self.random_contrast = RandomContrast(p=p, lower=lower, upper=upper)
        self.random_brightness = RandomBrightness(p=p, delta=delta)

    def __call__(self, image, boxes=None, labels=None):
        image, boxes, labels = self.random_contrast(image, boxes, labels)
        image, boxes, labels = self.random_brightness(image, boxes, labels)
        return image, boxes, labels


class RandomContrast(object):
    """随机调整对比度和亮度"""

    def __init__(self, p=0.5, lower=0.5, upper=1.5):
        self.p = p
        self.lower = lower
        self.upper = upper
        assert self.upper >= self.lower, "contrast upper must be >= lower."
        assert self.lower >= 0, "contrast lower must be non-negative."

    def __call__(self, image, boxes=None, labels=None):
        if random.random() < self.p:
            dtype = image.dtype
            image = np.asarray(image, dtype=np.float32)
            alpha = random.uniform(self.lower, self.upper)
            image *= alpha
            image = np.clip(image, 0, 255)
            image = np.asarray(image, dtype=dtype)
        return image, boxes, labels


class RandomBrightness(object):
    """随机调整亮度"""

    def __init__(self, p=0.5, delta=32):
        assert delta >= 0.0
        assert delta <= 255.0
        self.p = p
        self.delta = delta

    def __call__(self, image, boxes=None, labels=None):
        if random.random() < self.p:
            dtype = image.dtype
            image = np.asarray(image, dtype=np.float32)
            delta = random.uniform(-self.delta, self.delta)
            image += delta
            image = np.clip(image, 0, 255)
            image = np.asarray(image, dtype=dtype)
        return image, boxes, labels


class PhotometricDistort(object):
    def __init__(self):
        self.pd = [
            RandomContrast(),  # RGB
            ConvertColor(current="RGB", transform='HSV'),  # HSV
            RandomSaturation(),  # HSV
            RandomHue(),  # HSV
            ConvertColor(current='HSV', transform='RGB'),  # RGB
            RandomContrast()  # RGB
        ]
        self.rand_brightness = RandomBrightness()
        self.rand_light_noise = RandomLightingNoise()

    def __call__(self, image, boxes, labels):
        im = image.copy()
        im, boxes, labels = self.rand_brightness(im, boxes, labels)
        if random.randint(2):
            distort = Compose(self.pd[:-1])
        else:
            distort = Compose(self.pd[1:])
        im, boxes, labels = distort(im, boxes, labels)
        return self.rand_light_noise(im, boxes, labels)


class PortraitMode():
    def __init__(self):
        """
        竖屏模式
        """
        self.rot = RandomRot90(p=1.0)

    def __call__(self, img, boxes, labels):
        h, w, _ = img.shape
        if h <= w:
            # 如果h < w，是横屏图片,需要旋转为竖屏
            img, boxes, labels = self.rot(img, boxes, labels)
        return img, boxes, labels


class LandscapeMode():
    def __init__(self):
        """
        横屏模式
        """
        self.rot = RandomRot90(p=1.0)

    def __call__(self, img, boxes, labels):
        h, w, _ = img.shape
        if h >= w:
            # 如果h > w，是竖屏图片,需要旋转为横屏
            img, boxes, labels = self.rot(img, boxes, labels)
        return img, boxes, labels


class OrientationModel():
    def __init__(self, width, height):
        """
        根据长宽比，自适应旋转90°，获得对应的横屏或者竖屏的图像
        当width>height时，输出始终为获得横屏图像
        当width<height时，输出始终为获得竖屏图像
        :param width:
        :param height:
        """
        self.aspect_ratio = height / width
        self.portrait = PortraitMode()  # 竖屏模式
        self.landscape = LandscapeMode()  # 横屏模式

    def __call__(self, img, boxes, labels):
        if self.aspect_ratio >= 1.0:
            img, boxes, labels = self.portrait(img, boxes, labels)  # 竖屏模式
        else:
            img, boxes, labels = self.landscape(img, boxes, labels)  # 横屏模式
        return img, boxes, labels


class RandomOrientationModel():
    def __init__(self, p=0.5):
        """
        随机获得横屏或者竖屏的图像，相当于对输入图像随机旋转90°
        :param p: 随机获得横屏或者竖屏的图像的概率,值为0时，表示输出原始图像
        """
        self.p = p
        self.portrait = PortraitMode()  # 竖屏模式
        self.landscape = LandscapeMode()  # 横屏模式

    def __call__(self, img, boxes, labels):
        if random.random() < self.p:
            if random.random() < 0.5:
                img, boxes, labels = self.portrait(img, boxes, labels)  # 竖屏模式
            else:
                img, boxes, labels = self.landscape(img, boxes, labels)  # 横屏模式
        return img, boxes, labels


class RandomMosaic(object):
    def __init__(self, size, p=0.5, max_nums=6, samples=[2, 3, 4], flip=True, padding=False):
        """
        mosaic 数据增强
        :param size: 返回图片的大小[width,height]
        :param p: 随机概率
        :param max_nums: 如果单张图片检测框超过限制max_nums，则不进行mosaic数据增强，-1表示不限制
        :param samples: Mosaic需要的图片数据量[2,3,4]
                        2:使用2张图，随机进行上下拼接或者左右拼接
                        3:使用3张图，进行拼接
                        4:使用4张图，进行拼接
        """
        super(RandomMosaic, self).__init__()
        self.size = size
        self.max_nums = max_nums
        self.p = p
        self.padding = padding
        if isinstance(samples, numbers.Number):
            samples = [samples]
        assert 5 > max(samples), Exception("samples:{}".format(samples))
        assert 1 < min(samples), Exception("samples:{}".format(samples))
        self.samples = samples
        self.flip = flip
        self.data = []
        self.win_size = 4
        self.put_p = 1 / 16  # 随机缓存图片的概率，一般设置为1/batch-size

    def put_item(self, item):
        """添加item(行)到队列表中"""
        # item = self.resize_item(item, size=self.size, padding=False)
        if len(self.data) >= self.win_size:
            self.data.pop(0)
        self.data.append(item)
        for i in range(self.win_size - len(self.data)):
            item = copy.deepcopy(self.get_item(-1))
            self.data.append(item)

    def get_item(self, row=-1):
        """获取队列表中第index的item(行)数据"""
        item = self.data[row]
        return item

    def resize_item(self, item, size, padding=False):
        img, boxes, label = ScaleResize(size, 1.0, padding=padding)(item["img"], item["boxes"], item["label"])
        oitem = {"img": img, "boxes": boxes, "label": label}
        return oitem

    @staticmethod
    def random_hflip(image, boxes, labels, p=0.5):
        """随机水平翻转图片horizontal flip"""
        if random.random() < p:
            height, width, _ = image.shape
            image = image[:, ::-1, :]
            if len(boxes) > 0:
                boxes[:, [0, 2]] = width - boxes[:, [2, 0]]
        return image, boxes, labels

    def mosaic_for2samples(self, item_list, size, p=0.5):
        """需要二张图像进行拼接:item = {"img": img, "boxes": boxes, "label": label}"""
        assert len(item_list) == 2
        if random.random() < p:
            # 左右拼接
            item_list[0] = self.resize_item(item_list[0], size=[int(size[0] / 2), size[1]], padding=self.padding)
            item_list[1] = self.resize_item(item_list[1], size=[int(size[0] / 2), size[1]], padding=self.padding)
            h0, w0, d0 = item_list[0]["img"].shape
            img_new = np.hstack((item_list[0]["img"], (item_list[1]["img"])))
            boxes1 = item_list[1]["boxes"] + [w0, 0, w0, 0]
        else:
            # 上下拼接
            item_list[0] = self.resize_item(item_list[0], size=[size[0], int(size[1] / 2)], padding=self.padding)
            item_list[1] = self.resize_item(item_list[1], size=[size[0], int(size[1] / 2)], padding=self.padding)
            h0, w0, d0 = item_list[0]["img"].shape
            img_new = np.vstack((item_list[0]["img"], (item_list[1]["img"])))
            boxes1 = item_list[1]["boxes"] + [0, h0, 0, h0]
        dst_bboxes = np.vstack([item_list[0]["boxes"], boxes1])
        dst_labels = np.vstack([item_list[0]["label"], item_list[1]["label"]])
        item = {"img": img_new, "boxes": dst_bboxes, "label": dst_labels}
        return item

    def mosaic_for3samples(self, item_list, size):
        """需要三张图像进行拼接:item = {"img": img, "boxes": boxes, "label": label}"""
        assert len(item_list) == 3
        item_list[0] = self.resize_item(item_list[0], size=[int(size[0] / 2), size[1]], padding=self.padding)
        item_list[1] = self.resize_item(item_list[1], size=[int(size[0] / 2), int(size[1] / 2)], padding=self.padding)
        item_list[2] = self.resize_item(item_list[2], size=[int(size[0] / 2), int(size[1] / 2)], padding=self.padding)
        # 垂直方向拼接
        img_12 = np.vstack((item_list[1]["img"], item_list[2]["img"]))
        # 水平方向拼接
        h0, w0, d0 = item_list[0]["img"].shape
        img_new = np.hstack((item_list[0]["img"], img_12))
        boxes1 = item_list[1]["boxes"] + [w0, 0, w0, 0]
        boxes2 = item_list[2]["boxes"] + [w0, int(h0 / 2), w0, int(h0 / 2)]
        dst_bboxes = np.vstack([item_list[0]["boxes"], boxes1, boxes2])
        dst_labels = np.vstack([item_list[0]["label"], item_list[1]["label"],
                                item_list[2]["label"]])
        item = {"img": img_new, "boxes": dst_bboxes, "label": dst_labels}
        return item

    def mosaic_for4samples(self, item_list, size=None):
        """需要四张图像进行拼接:item = {"img": img, "boxes": boxes, "label": label}"""
        assert len(item_list) == 4
        item_list[0] = self.resize_item(item_list[0], size=[int(size[0] / 2), int(size[1] / 2)], padding=self.padding)
        item_list[1] = self.resize_item(item_list[1], size=[int(size[0] / 2), int(size[1] / 2)], padding=self.padding)
        item_list[2] = self.resize_item(item_list[2], size=[int(size[0] / 2), int(size[1] / 2)], padding=self.padding)
        item_list[3] = self.resize_item(item_list[3], size=[int(size[0] / 2), int(size[1] / 2)], padding=self.padding)
        # 水平方向拼接
        img_01 = np.hstack((item_list[0]["img"], item_list[1]["img"]))
        img_23 = np.hstack((item_list[2]["img"], item_list[3]["img"]))
        # 垂直方向拼接
        img_new = np.vstack((img_01, img_23))
        h0, w0, d0 = item_list[0]["img"].shape
        boxes1 = item_list[1]["boxes"] + [w0, 0, w0, 0]
        boxes2 = item_list[2]["boxes"] + [0, h0, 0, h0]
        boxes3 = item_list[3]["boxes"] + [w0, h0, w0, h0]
        dst_bboxes = np.vstack([item_list[0]["boxes"], boxes1, boxes2, boxes3])
        dst_labels = np.vstack([item_list[0]["label"], item_list[1]["label"],
                                item_list[2]["label"], item_list[3]["label"]])
        item = {"img": img_new, "boxes": dst_bboxes, "label": dst_labels}
        return item

    def mosaic(self, item, samples):
        item0 = copy.deepcopy(item)
        if samples == 2:
            item1 = copy.deepcopy(self.get_item(0))
            item_list = [item0, item1]
            out_item = self.mosaic_for2samples(item_list, size=self.size)
        elif samples == 3:
            item1 = copy.deepcopy(self.get_item(0))
            item2 = copy.deepcopy(self.get_item(1))
            item_list = [item0, item1, item2]
            out_item = self.mosaic_for3samples(item_list, size=self.size)
        elif samples == 4:
            item1 = copy.deepcopy(self.get_item(0))
            item2 = copy.deepcopy(self.get_item(1))
            item3 = copy.deepcopy(self.get_item(2))
            item_list = [item0, item1, item2, item3]
            out_item = self.mosaic_for4samples(item_list, size=self.size)
        else:
            raise Exception("Error:samples:{}".format(self.samples))
        return out_item

    def __call__(self, img, boxes, label):
        item = {"img": img, "boxes": boxes, "label": label}
        if len(boxes) > self.max_nums and self.max_nums > 0:
            return img, boxes, label
        if len(boxes) > 0:
            # 随机保留训练图片，用于mosaic
            if random.random() < self.put_p or len(self.data) == 0:
                self.put_item(item)
            # 随机生成mosaic图片
            if random.random() < self.p:
                samples = np.random.choice(self.samples)
                out_item = self.mosaic(item, samples)
                img, boxes, label = out_item["img"], out_item["boxes"], out_item["label"]
                if self.flip:
                    img, boxes, label = self.random_hflip(img, boxes, label)
        boxes = np.asarray(boxes, dtype=np.float32)
        return img, boxes, label


class RandomMixup(object):
    def __init__(self, size, p=0.5, flip=True, lambd=0.5):
        """
        Mixup数据增强
        :param size: 返回图片的大小[width,height]
        :param p: 随机概率
        :param flip: 随机翻转
        :param lambd: Mixup混合程度，默认0.5,设置为None时，则采用狄利克雷分布np.random.beta
        """
        super(RandomMixup, self).__init__()
        self.size = size
        self.p = p
        self.data = []
        self.flip = flip
        self.win_size = 4
        self.put_p = 1 / 20
        self.lambd = lambd

    def put_item(self, item):
        """添加item(行)到队列表中"""
        item = self.resize_item(item, size=self.size)
        if len(self.data) >= self.win_size:
            self.data.pop(0)
        self.data.append(item)
        for i in range(self.win_size - len(self.data)):
            item = self.get_item(-1).copy()
            self.data.append(item)

    def get_item(self, row=-1):
        """获取队列表中第index的item(行)数据"""
        item = self.data[row]
        return item

    @staticmethod
    def resize_item(item, size):
        img, boxes, label = Resize(size=size)(item["img"], item["boxes"], item["label"])
        item["img"] = img
        item["boxes"] = boxes
        item["label"] = label
        return item

    @staticmethod
    def random_hflip(image, boxes, labels, p=0.5):
        """随机水平翻转图片horizontal flip"""
        if random.random() < p:
            height, width, _ = image.shape
            image = image[:, ::-1, :]
            if len(boxes) > 0:
                boxes[:, [0, 2]] = width - boxes[:, [2, 0]]
        return image, boxes, labels

    def mixup(self, item0, item1, lambd=None):
        """Y
        OLOv5:utils/augmentations.py
        狄利克雷分布np.random.beta: 概率接近0.5
        """
        # Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf
        # r = np.random.beta(32.0, 32.0)  # mixup ratio, alpha=beta=32.0
        # im = (im * r + im2 * (1 - r)).astype(np.uint8)
        # labels = np.concatenate((labels, labels2), 0)
        if not lambd:
            lambd = np.random.beta(32.0, 32.0)  # mixup ratio, alpha=beta=32.0
        image = item0["img"] * lambd + item1["img"] * (1 - lambd)
        item0["img"] = np.asarray(image, dtype=np.uint8)
        item0["boxes"] = np.concatenate((item0["boxes"], item1["boxes"]), axis=0)
        item0["label"] = np.concatenate((item0["label"], item1["label"]), axis=0)
        return item0

    def __call__(self, img, boxes, label):
        item0 = {"img": img, "boxes": boxes, "label": label}
        if len(boxes) > 0:
            # 随机保留训练图片，用于Mixup
            if random.random() < self.put_p or len(self.data) == 0:
                self.put_item(item0)
            # 随机生成Mixup图片
            if random.random() < self.p:
                item0 = self.resize_item(item0, size=self.size)
                item1 = self.get_item(0).copy()
                out_item = self.mixup(item0, item1, lambd=self.lambd)
                img, boxes, label = out_item["img"], out_item["boxes"], out_item["label"]
                if self.flip:
                    img, boxes, label = self.random_hflip(img, boxes, label)
        boxes = np.asarray(boxes, dtype=np.float32)
        return img, boxes, label


class RandomBoxesPaste(object):
    """ 实现随机背景贴图"""

    def __init__(self, p=0.5, bg_dir="bg_image/", use_rgb=True):
        """
        实现随机背景贴图，适用于检测对象之间没有overlap关系的目标检测数据增强
        使用条件：
        (1)检测对象之间没有overlap关系,如：face和person之间boxes经常存在overlap关系，
           随机贴图容易出现face的box缺漏的问题
        (2)背景图不能含有目标检测的对象，避免背景的干扰
        (3)对于单目标检测，都可以使用该方法，实现数据增强
        (4)支持最多2个boxes的随机贴图，如果输入boxes多余2个，则自动丢弃多余的boxes
        :param p: 实现随机背景贴图的概率
        :param bg_dir: 背景图库，PS：背景图不能含有目标检测的对象，避免背景的干扰
        """
        from detection.utils import file_utils
        self.p = p
        self.max_scale = 2 / 5
        self.scale = [1.2, 1.2]
        self.use_rgb = use_rgb
        self.bg_dir = bg_dir
        self.bg_nums = 0
        if self.bg_dir:
            self.bg_image_list = file_utils.get_files_lists(bg_dir, subname="")
            self.bg_nums = len(self.bg_image_list)
        print("use RandomBoxesPaste:p:{},bg_dir:{},use_rgb:{},have {} images".format(p, bg_dir, use_rgb,
                                                                                     self.bg_nums))

    def random_read_bg_image(self, use_rgb=False, crop_rate=0.5):
        index = int(np.random.uniform(0, self.bg_nums))
        image_path = self.bg_image_list[index]
        # image_path = self.bg_image_list[0]
        image = cv2.imread(image_path)
        if use_rgb:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        h_img, w_img, _ = image.shape
        xmin, ymin, xmax, ymax = self.extend_bboxes([0, 0, w_img, h_img], [0.8, 0.8])
        crop_xmin = int(random.uniform(0, xmin * crop_rate))
        crop_ymin = int(random.uniform(0, ymin * crop_rate))
        crop_xmax = int(min(w_img, random.uniform(w_img - xmax * crop_rate, w_img)))
        crop_ymax = int(min(h_img, random.uniform(h_img - ymax * crop_rate, h_img)))
        if random.random() < 0.5:
            image = image[:, ::-1, :]
        image = image[crop_ymin: crop_ymax, crop_xmin: crop_xmax]
        return image

    @staticmethod
    def extend_bboxes(box, scale=[1.0, 1.0]):
        """
        :param box: [xmin, ymin, xmax, ymax]
        :param scale: [sx,sy]==>(W,H)
        :return:
        """
        sx = scale[0]
        sy = scale[1]
        xmin, ymin, xmax, ymax = box[:4]
        cx = (xmin + xmax) / 2
        cy = (ymin + ymax) / 2

        ex_w = (xmax - xmin) * sx
        ex_h = (ymax - ymin) * sy
        ex_xmin = cx - 0.5 * ex_w
        ex_ymin = cy - 0.5 * ex_h
        ex_xmax = ex_xmin + ex_w
        ex_ymax = ex_ymin + ex_h
        ex_box = [ex_xmin, ex_ymin, ex_xmax, ex_ymax]
        return ex_box

    @staticmethod
    def cv_paste_image(im, mask, boxes, start_point=(0, 0)):
        """
        :param im:
        :param start_point:
        :param mask:
        :return:
        """
        xim, ymin = start_point
        shape = mask.shape  # h, w, d
        im[ymin:(ymin + shape[0]), xim:(xim + shape[1])] = mask
        boxes = boxes + [start_point[0], start_point[1], start_point[0], start_point[1]]
        return im, boxes

    def center2box(self, center):
        cx, cy, w, h = center
        box = [int(cx - w / 2), int(cy - h / 2), int(cx + w / 2), int(cy + h / 2)]
        return box

    @staticmethod
    def pil_paste_image(im, mask, boxes, start_point=(0, 0)):
        """
        :param im:
        :param mask:
        :param start_point: tupe
        :return:
        """
        assert isinstance(start_point, tuple)
        out = Image.fromarray(im)
        mask = Image.fromarray(mask)
        out.paste(mask, start_point)
        boxes = boxes + [start_point[0], start_point[1], start_point[0], start_point[1]]
        return np.asarray(out), boxes

    @staticmethod
    def get_random_point(w, h, edge_wh=[0, 0]):
        """
        获得随机点
        :param w:
        :param h:
        :param edge_wh: 右小角的边缘宽度
        :return:
        """
        dst_w = int(np.random.uniform(0, w - edge_wh[0]))
        dst_h = int(np.random.uniform(0, h - edge_wh[1]))
        return [dst_w, dst_h]

    def get_crop_roi(self, img, box, edge_wh_scale=[1.0, 1.0]):
        """
        获得裁剪区域ROI
        :param img:
        :param box:
        :param edge_wh_scale: 裁剪区域ROI外扩大小
        :return:
        """
        h_img, w_img, _ = img.shape
        xmin, ymin, xmax, ymax = self.extend_bboxes(box, edge_wh_scale)
        xmin = max(0, xmin)
        ymin = max(0, ymin)
        xmax = min(w_img, xmax)
        ymax = min(h_img, ymax)
        roi = img[int(ymin):int(ymax), int(xmin):int(xmax)]
        box = box - [xmin, ymin, xmin, ymin]
        return roi, box

    def check_distance(self, point0, point1, dis_th):
        """
        检测点与点之间的距离关系
        :param point0:
        :param point1:
        :param dis_th: 距离阈值
        :return: True:存在overlap，False:不存在overlap
        """
        r = abs(point0[0] - point1[0]) < dis_th and abs(point0[1] - point1[1]) < dis_th
        return r

    def check_point(self, point, max_wh, roi_wh):
        rw = point[0] + roi_wh[0] < max_wh[0] and point[0] >= 0
        rh = point[1] + roi_wh[1] < max_wh[1] and point[1] >= 0
        return rw * rh

    def random_paste(self, img, boxes, labels, bg_image):
        sh, sw, _ = img.shape
        bh, bw, d = bg_image.shape
        ratio = [sw / bw, sh / bh]
        ratio = max(ratio)
        if ratio > 1:
            bg_image = cv2.resize(bg_image, dsize=(math.ceil(bw * ratio), math.ceil(bh * ratio)))
        bg_image = bg_image[0: sh, 0: sw]
        start_point = [0, 0]
        bg_boxes = []
        bg_labels = []
        max_wh = 0
        # 目前支持最多2个boxes的随机贴图
        for i in range(min(len(boxes), 2)):
            box = boxes[i, :]
            label = labels[i]
            roi, box = self.get_crop_roi(img, box)
            roi_h, roi_w, _ = roi.shape
            # dis_th：贴图之间的间距，避免overlap
            dis_th = max(roi_w, roi_h, max_wh)
            for k in range(10):
                point = self.get_random_point(sw, sh, edge_wh=[roi_w, roi_h])
                # print("(i,k)=({},{}),point0：{}，point1：{}".format(i, k, point, start_point))
                if not self.check_point(point, max_wh=(sw, sh), roi_wh=[roi_w, roi_h]):
                    continue
                elif not self.check_distance(point, start_point, dis_th=dis_th * 0.9):
                    start_point = point
                    bg_image, box = self.pil_paste_image(bg_image, roi, box, start_point=tuple(start_point))
                    bg_boxes.append(box)
                    bg_labels.append(label)
                    break
            max_wh = dis_th
        if len(bg_boxes) == 0:
            bg_boxes = boxes
            bg_labels = labels
            bg_image = img
        bg_boxes = np.asarray(bg_boxes)
        bg_labels = np.asarray(bg_labels)
        return bg_image, bg_boxes, bg_labels

    def __call__(self, img, boxes, labels):
        """
        Args:
            img (numpy Image): Image to be cropped.
        Returns:
            PIL Image: Cropped image.
        """
        if not self.bg_dir:
            return img, boxes, labels
        if random.random() < self.p:
            bg_image = self.random_read_bg_image(use_rgb=self.use_rgb)
            img, boxes, labels = self.random_paste(img, boxes, labels, bg_image)
        return img, boxes, labels


def show_image(image, bboxes, labels, normal=False, transpose=False):
    """
    :param image:
    :param targets_t:
                bboxes = targets[idx][:, :4].data
                keypoints = targets[idx][:, 4:14].data
                labels = targets[idx][:, -1].data
    :return:
    """
    import numpy as np
    from detection.utils import image_utils
    image = np.asarray(image)
    bboxes = np.asarray(bboxes)
    labels = np.asarray(labels)
    print("bboxes:{}".format(bboxes))
    print("labels:{}".format(labels))
    if transpose:
        image = image_utils.untranspose(image)
    h, w, _ = image.shape
    landms_scale = np.asarray([w, h] * 5)
    bboxes_scale = np.asarray([w, h] * 2)
    if normal:
        bboxes = bboxes * bboxes_scale
    # tmp_image = image_processing.untranspose(tmp_image)
    tmp_image = image_utils.convert_color_space(image, colorSpace="BGR")
    tmp_image = image_utils.draw_image_bboxes_text(tmp_image, bboxes, labels)
    image_utils.cv_show_image("train", tmp_image)


def demo_for_augment():
    from detection.utils import image_utils
    bg_dir = "/home/dm/data3/dataset/finger_keypoint/bug/JPEGImages"
    input_size = [320, 320]
    image_path = "test.jpg"
    boxes = [[98, 42, 160, 100], [244, 260, 297, 332], [98 + 50, 42 + 50, 160 + 50, 100 + 50]]
    labels = [1, 2, 3]
    image = image_utils.read_image(image_path)
    boxes = np.asarray(boxes)
    labels = np.asarray(labels)
    augment = Compose([
        # augment_bbox_landm.ProtectBoxes(),
        # augment_bbox.RandomRot90(),  # 随机横屏和竖屏
        # RandomRotation(degrees=15),
        ProtectBoxes(norm=False),
        # RandomHorizontalFlip(flip_index=flip_index),
        RandomBoxesPaste(bg_dir=bg_dir),
        # RandomVerticalFlip(),
        # RandomCrop(),
        # RandomCropLarge(size=input_size),
        # RandomContrastBrightness(),
        RandomAffineResizePadding(degrees=15, output_size=input_size),
        RandomMosaic(size=input_size),
        # RandomAffineResizePadding(degrees=15),
        # ResizePadding(input_size),
        # RandomResizePadding(input_size, p=0.5),
        # Resize(input_size),
        # ResizePadding(input_size),
        RandomColorJitter(),
        # SwapChannels(),
    ])

    for i in range(1000):
        print("===" * 10)
        dst_image, dst_boxes, dst_labels = augment(image.copy(), boxes.copy(), labels.copy())
        show_image(dst_image, dst_boxes, dst_labels, normal=False, transpose=False)


def demo_for_rotation():
    from detection.utils import image_utils
    input_size = [320, 320]
    image_path = "test.jpg"
    images = image_utils.read_image(image_path)
    bboxes = [[98, 42, 160, 100], [98 + 50, 42 + 50, 160 + 50, 100 + 50], [244, 260, 297, 332], [124, 126, 125, 135]]
    labels = np.asarray([1, 2, 3, 4])
    bboxes = np.asarray(bboxes)

    # augment = RandomRotation(degrees=15, p=1.0)
    augment = RandomAffineResizePadding(degrees=15, output_size=input_size, p=1.0)
    # augment = RandomAffine()
    augment = RandomMosaic(size=input_size, p=0.9, samples=[2, 3, 4])
    augment = ResizePadding(size=input_size)
    # augment = Resize(size=input_size)
    # 顺时针旋转90度
    # image_processing.show_image_boxes("src", images, bboxes, color=(0, 255, 0), waitKey=10)
    for angle in range(360 * 10):
        # angle = 30
        # augment.set_degrees(angle)
        dst_image, dst_boxes, dst_labels = augment(images.copy(), bboxes.copy(), labels.copy())
        # dst_image, dst_boxes, dst_labels = augment.inverse(dst_image.copy(), dst_boxes.copy(), dst_labels.copy())
        dst_image, dst_boxes, dst_labels = augment(dst_image.copy(), dst_boxes.copy(), dst_labels.copy())
        print("==" * 10)
        print("angle:{}".format(angle))
        print("shape:{},bboxes     ：{}".format(bboxes.shape, bboxes))
        print("dst_image:{},dst_boxes：{},dst_labels:{}".format(dst_image.shape, dst_boxes, dst_labels))
        show_image(dst_image, dst_boxes, dst_labels, normal=False, transpose=False)


if __name__ == "__main__":
    demo_for_augment()
    # demo_for_rotation()
